import torch
import torch.utils.data as Data
from dataset import BookUserInfo, TrainDataset
from evaluate import Evaluator
from model import NewFMModel
from tqdm import tqdm
import random
import numpy as np
import os

os.environ['CUDA_VISIBLE_DEVICES'] = '2'

def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True

setup_seed(1)

hyper_params = {
    'prerec_result': 'model_dat/prerec_result.pkl',
    'data_path': 'datasets/',
    'max_title_len': 10,
    'feature_dim': 32,
    'fm_dim': 32,
    'batch_size': 512,
    'eval_batch_size': 1024,
    'device': 'cuda',
    'epochs': 1000
}

def train(book_user_info: BookUserInfo):
    # Load Data
    print('train: preparing...')
    dataset = TrainDataset(hyper_params, book_user_info)
    loader = Data.DataLoader(dataset, batch_size=hyper_params['batch_size'], shuffle=True)

    # Define Model
    model = NewFMModel(hyper_params).to(hyper_params['device'])

    # Define Optimizer
    optimizer = torch.optim.AdamW(model.parameters())

    # Define Evaluator
    evaluator = Evaluator(hyper_params, book_user_info)

    # Start training
    print('train: start training...')

    for epoch in range(hyper_params['epochs']):
        print(f'train: Epoch {epoch}')
        total_loss = 0.0
        total_step = 0

        #t = tqdm(loader)
        for data in loader:
            user_id, school, grade, pos_bid, pos_title, pos_type, neg_bid, neg_title, neg_type = [i.to(hyper_params['device']) for i in data]

            pos_score = model(user_id, school, grade, pos_bid, pos_title, pos_type)
            neg_score = model(user_id, school, grade, neg_bid, neg_title, neg_type)

            loss = -torch.log(torch.sigmoid(pos_score-neg_score)).mean()
            total_loss += loss.item()
            total_step += 1

            # t.set_description(f'loss: {loss.item()}')

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

        print(f'Epoch: {epoch}, Loss: {total_loss / total_step}')
        
        if epoch % 5 == 0:
            evaluator.evaluate(model)


if __name__ == '__main__':
    book_user_info = BookUserInfo(hyper_params)
    hyper_params = book_user_info.update_hyper_params()

    train(book_user_info)
